use std::rc::Rc;
use itertools::Itertools;
use num_traits::Num;
use crate::autograd::function::{ForwardArgs, Function};
use crate::autograd::function_ctx::FunctionCtx;
use crate::Tensor;

pub struct PermuteBackward<T> {
    ctx: FunctionCtx<T>,
}

impl<T> Function<T> for PermuteBackward<T>
    where
        T: Copy + Num + 'static,
{
    fn new(ctx: FunctionCtx<T>) -> Self {
        Self { ctx }
    }

    fn ctx(&self) -> &FunctionCtx<T> {
        &self.ctx
    }

    fn forward(ctx: &mut FunctionCtx<T>, args: ForwardArgs<T>) -> Tensor<T> {
        if let ForwardArgs::TensorOther(tensor, axis) = args {
            ctx.save_others([axis].into());
            let cell = (*tensor.ptr).borrow();

            assert_eq!(
                axis.len(), cell.shape.len(),
                "number of dimensions in the tensor input does not match \
                the length of the desired ordering of dimensions \
                i.e. input.dim() = {:?} is not equal to len(dims) = {:?}",
                cell.shape.len(), axis.len());

            let shape = axis.iter()
                .map(|i| cell.shape[*i])
                .collect();

            let strides = axis.iter()
                .map(|i| cell.strides[*i])
                .collect();

            return (Rc::clone(&cell.data), cell.offset, shape, strides).into()
        }
        unreachable!()
    }

    fn backward(&self, grad_output: Tensor<T>) -> Vec<Option<Tensor<T>>> {
        let axis = self.ctx.others()[0];
        let default_axis = (0..axis.len()).collect_vec();

        let axis = axis.iter()
            .zip(default_axis)
            .sorted_by_key(|x| x.0)
            .map(|x| x.1)
            .collect_vec();

        [grad_output.permute(&axis).into()].into()
    }
}
